import tkinter as tk
from tkinter import messagebox
import sqlite3
from mnist_recognize import RecognizeApp

# 数据库操作函数
def register_user(username, password):
    conn = sqlite3.connect('user_data.db')
    c = conn.cursor()
    c.execute('INSERT INTO users (username, password) VALUES (?, ?)', (username, password))
    conn.commit()
    conn.close()

def validate_login(username, password):
    conn = sqlite3.connect('user_data.db')
    c = conn.cursor()
    c.execute('SELECT * FROM users WHERE username = ? AND password = ?', (username, password))
    result = c.fetchone()
    conn.close()
    return result

# 登录和注册界面
class LoginApp(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title('用户登录')
        self.geometry('300x200')

        self.label_username = tk.Label(self, text='用户名')
        self.label_username.pack(pady=5)
        self.entry_username = tk.Entry(self)
        self.entry_username.pack(pady=5)

        self.label_password = tk.Label(self, text='密码')
        self.label_password.pack(pady=5)
        self.entry_password = tk.Entry(self, show='*')
        self.entry_password.pack(pady=5)

        self.button_login = tk.Button(self, text='登录', command=self.login)
        self.button_login.pack(pady=5)
        
        self.button_register = tk.Button(self, text='注册', command=self.register)
        self.button_register.pack(pady=5)

    def login(self):
        username = self.entry_username.get()
        password = self.entry_password.get()
        if validate_login(username, password):
            messagebox.showinfo('登录成功', '欢迎，' + username)
            self.open_recognize_app()
        else:
            messagebox.showerror('登录失败', '用户名或密码错误')

    def register(self):
        username = self.entry_username.get()
        password = self.entry_password.get()
        if len(username) < 6 or len(password) < 6:
            messagebox.showerror('注册失败', '用户名和密码长度必须不少于6位')
        else:
            register_user(username, password)
            messagebox.showinfo('注册成功', '请重新登录')

    def open_recognize_app(self):
        self.withdraw()
        recognize_app = RecognizeApp(self)
        recognize_app.mainloop()

if __name__ == "__main__":
    app = LoginApp()
    app.mainloop()